{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5f29237f",
   "metadata": {},
   "source": [
    "# SFT_GRPO_2\n",
    ": Reward functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aab95e1a",
   "metadata": {},
   "source": [
    "Start by importing dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a26ad554-af62-4be1-af64-6532cc6c4caf",
   "metadata": {
    "height": 98
   },
   "outputs": [],
   "source": [
    "from utils import *\n",
    "import numpy as np\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40bf03b2-8d85-4f45-a178-bd5c32b52e0b",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b71ee2bf",
   "metadata": {},
   "source": [
    "Create a predibase deployment that you'll use to call the models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0103d4ab-ba27-4469-84e1-270aa18728c8",
   "metadata": {
    "height": 62
   },
   "outputs": [],
   "source": [
    "# Uncomment the line below if running in your own environment - the deployment is already setup for you here\n",
    "# create_deployment()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11bfd21c-9f62-4e75-bb50-0a814d772ac7",
   "metadata": {},
   "source": [
    "Specify the model to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aebfa2b-ea17-4d93-9f17-2505cd3ce12a",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "model_id = \"Qwen/Qwen2.5-7B-Instruct\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c04eeb",
   "metadata": {},
   "source": [
    "## Define a simple reward function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a74338bd-6515-4e1e-820a-3a3881aac9dc",
   "metadata": {
    "height": 98
   },
   "outputs": [],
   "source": [
    "def wordle_reward(guess: str, secret_word: str) -> int:\n",
    "    if guess.upper() == secret_word.upper():\n",
    "        return 1   # correct guess\n",
    "    else:\n",
    "        return 0   # incorrect guess"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "838e186a",
   "metadata": {},
   "source": [
    "Define a secret word and get feedback on past guesses, then score the guesses using the reward function above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "822c372c-daa3-430d-9cbb-d71946a9d617",
   "metadata": {
    "height": 164
   },
   "outputs": [],
   "source": [
    "secret_word = \"POUND\"\n",
    "\n",
    "past_guesses = [\n",
    "    GuessWithFeedback.from_secret(guess=\"CRANE\", secret=secret_word),\n",
    "    GuessWithFeedback.from_secret(guess=\"BLOND\", secret=secret_word),\n",
    "    GuessWithFeedback.from_secret(guess=\"FOUND\", secret=secret_word),\n",
    "]\n",
    "past_guesses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce53f2f2-0301-4ec4-a83d-e134dd0e6ec5",
   "metadata": {
    "height": 98
   },
   "outputs": [],
   "source": [
    "response = generate(get_messages(past_guesses))[0]\n",
    "guess = extract_guess(response)\n",
    "reward = wordle_reward(guess, secret_word)\n",
    "\n",
    "print(f\"Guessed Word: {guess} -> Reward: {reward}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "823ec508",
   "metadata": {},
   "source": [
    "## Using rewards to calculate advantages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75c6266b-e5e1-4189-8e96-2a01abe78ba5",
   "metadata": {
    "height": 283
   },
   "outputs": [],
   "source": [
    "def compute_advantages(rewards: list):\n",
    "    rewards = np.array(rewards)\n",
    "    \n",
    "    # Compute the mean and standard deviation of the rewards\n",
    "    mean_reward = np.mean(rewards)\n",
    "    std_reward = np.std(rewards)\n",
    "\n",
    "    # Avoid division by zero in case of zero variance (typically happens when all rewards are 0)\n",
    "    # Note: In the GRPO implementation, we add 1e-4 to the std_reward to avoid division by zero\n",
    "    if std_reward == 0:\n",
    "        return [0] * len(rewards)\n",
    "\n",
    "    # Divide by stddev of rewards to normalize range to 0\n",
    "    advantages = (rewards - mean_reward) / std_reward\n",
    "    return advantages.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8109424-cca3-49f7-9a59-117a1b060687",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "rewards = [0.0, 0.2, 0.4, 0.5, 0.5, 0.6, 0.8, 1.0]\n",
    "compute_advantages(rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "794ff096-ce27-4df7-acd2-79a74f6ae8ba",
   "metadata": {
    "height": 96
   },
   "outputs": [],
   "source": [
    "def render_guess_table(response, reward_fn):\n",
    "    guesses = [extract_guess(guess) for guess in response]\n",
    "    rewards = [reward_fn(guess, secret_word) for guess in guesses]\n",
    "    print_guesses_table(guesses, rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2bf5d92-76da-4b86-b9c7-08413df11dff",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "print(f\"Secret: {secret_word}\")\n",
    "response = generate(get_messages(past_guesses), num_guesses=8)\n",
    "render_guess_table(response, wordle_reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c8047a",
   "metadata": {},
   "source": [
    "## Update the reward function to give partial credit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00b73314-031c-493b-904e-86a856f0671c",
   "metadata": {
    "height": 334
   },
   "outputs": [],
   "source": [
    "def wordle_reward_partial_credit(guess: str, secret_word: str) -> float:\n",
    "    if len(guess) != len(secret_word):\n",
    "        # no reward for having the wrong number of letters\n",
    "        return 0.0\n",
    "    \n",
    "    valid_letters = set(secret_word)\n",
    "    reward = 0.0\n",
    "    for letter, secret_letter in zip(guess, secret_word):\n",
    "        if letter == secret_letter:\n",
    "            # right letter, right location\n",
    "            reward += 0.2\n",
    "        elif letter in valid_letters:\n",
    "            # right letter, wrong location\n",
    "            reward += 0.1\n",
    "        else:\n",
    "            # no reward\n",
    "            pass\n",
    "    return reward"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51d84113",
   "metadata": {},
   "source": [
    "Try scoring a set of responses using updated reward function. Start by setting <b>temperature = 0</b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b40d00d-9d3a-4aff-9fcc-cb3abbd5f2ce",
   "metadata": {
    "height": 79
   },
   "outputs": [],
   "source": [
    "print(f\"Secret: {secret_word}\")\n",
    "response = generate(get_messages(past_guesses), num_guesses=8, temperature=0)\n",
    "render_guess_table(response, wordle_reward_partial_credit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bde7a8f0",
   "metadata": {},
   "source": [
    "Now set temperature to a high value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "386c8a7d-d683-4286-b971-21dc6cafdd8c",
   "metadata": {
    "height": 79
   },
   "outputs": [],
   "source": [
    "print(f\"Secret: {secret_word}\")\n",
    "response = generate(get_messages(past_guesses), num_guesses=8, temperature=1.3)\n",
    "render_guess_table(response, wordle_reward_partial_credit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6302bf",
   "metadata": {},
   "source": [
    "Lastly, set temperature to a moderate value of 0.7:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "307607c3-39db-4104-b745-408320b8b77a",
   "metadata": {
    "height": 79
   },
   "outputs": [],
   "source": [
    "print(f\"Secret: {secret_word}\")\n",
    "response = generate(get_messages(past_guesses), num_guesses=8, temperature=0.7)\n",
    "render_guess_table(response, wordle_reward_partial_credit)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
